from io import BytesIO, TextIOWrapper

from textfsm import TextFSM


class ZetTextFsm(object):

    def __init__(self, template=None, template_str=None):
        if template:
            self.fsm = TextFSM(open(template))
        elif template_str:
            bio = BytesIO(bytes(template_str, encoding='utf8'))  # 像文件对象一样操作
            text_iofile = TextIOWrapper(bio, encoding='utf8')
            self.fsm = TextFSM(text_iofile)
        else:
            raise ValueError('template 或者 template_str 至少一个非空')

    def parse_2_structured_data(self, text, multi=True):
        ''':arg
        将文本解析成为结构化数据，如果multi是多个，返回的是列表，否则返回dict
        如果想返回一个必须是列表长度为1 且multi=false
        '''
        objs = self.fsm.ParseTextToDicts(text)
        if objs:
            if len(objs) == 1 and not multi:
                return objs[0]
            else:
                return objs

        return None


if __name__ == '__main__':
    with open('show_version.textfsm', 'r', encoding='utf8') as f:
        templ_str = f.read()

    with open('show_version.log', 'r', encoding='utf8') as f:
        dev_text = f.read()

    # 用字符串构建zet textfsm
    fsm = ZetTextFsm(template_str=templ_str)
    data = fsm.parse_2_structured_data(dev_text, multi=False)
    print(data)
    # 用文件名构建zet textfsm
    fsm = ZetTextFsm(template='show_version.textfsm')
    data = fsm.parse_2_structured_data(dev_text, multi=False)
    print(data)
